import torch
import torch.nn as nn
from torchinfo import summary

# 3D卷积块：包含两个连续的3D卷积层，每个卷积层后接BatchNorm和ReLU激活
class Conv3D_Block(nn.Module):
    def __init__(self, in_feats, out_feats, kernel_size=3, padding=1):
        super().__init__()
        # 定义双卷积结构：两个连续的3D卷积+BN+ReLU
        self.conv = nn.Sequential(nn.Conv3d(in_channels=in_feats, out_channels=out_feats, kernel_size=kernel_size, padding=padding, bias=True),
                                  nn.BatchNorm3d(out_feats),
                                  nn.ReLU(inplace=True),
                                  nn.Conv3d(in_channels=out_feats, out_channels=out_feats, kernel_size=kernel_size, padding=padding, bias=True),
                                  nn.BatchNorm3d(out_feats),
                                  nn.ReLU(inplace=True))

    def forward(self, x):
        return self.conv(x)

# 上采样块：先进行三线性插值上采样，然后进行3D卷积减少通道数
class Up_Block(nn.Module):
    def __init__(self, in_feats):
        super().__init__()
        # 三线性插值上采样，尺寸扩大2倍
        self.up = nn.Upsample(scale_factor=2, mode="trilinear", align_corners=True)
        # 定义一个3D卷积层，用于减少通道数为原来的一半
        self.conv = nn.Conv3d(in_channels=in_feats, out_channels=int(in_feats / 2), kernel_size=3, stride=1, padding=1, bias=True)

    def forward(self, x):
        return self.conv(self.up(x))


# 3D U-Net网络实现
class UNet3D(nn.Module):
    def __init__(self, input_ch=1, output_ch=14, init_feats=16):
        super().__init__()
        # 编码器部分 - 特征提取
        self.conv_block_1 = Conv3D_Block(in_feats=input_ch, out_feats=init_feats)
        self.conv_block_2 = Conv3D_Block(in_feats=init_feats, out_feats=init_feats * 2)
        self.conv_block_3 = Conv3D_Block(in_feats=init_feats * 2, out_feats=init_feats * 4)
        self.conv_block_4 = Conv3D_Block(in_feats=init_feats * 4, out_feats=init_feats * 8)
        self.conv_block_5 = Conv3D_Block(in_feats=init_feats * 8, out_feats=init_feats * 16)
        # 下采样操作
        self.max_pool = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2))
        # 解码器部分 - 特征恢复与上采样
        self.up_block_1 = Up_Block(in_feats=init_feats * 16)
        self.up_block_2 = Up_Block(in_feats=init_feats * 8)
        self.up_block_3 = Up_Block(in_feats=init_feats * 4)
        self.up_block_4 = Up_Block(in_feats=init_feats * 2)
        # 结合跳跃连接后的卷积处理
        self.conv_block_6 = Conv3D_Block(in_feats=init_feats * 16, out_feats=init_feats * 8)
        self.conv_block_7 = Conv3D_Block(in_feats=init_feats * 8, out_feats=init_feats * 4)
        self.conv_block_8 = Conv3D_Block(in_feats=init_feats * 4, out_feats=init_feats * 2)
        self.conv_block_9 = Conv3D_Block(in_feats=init_feats * 2, out_feats=init_feats)
        # 最终的1x1卷积，用于生成输出分割图
        self.out = nn.Conv3d(in_channels=init_feats, out_channels=output_ch, kernel_size=1, stride=1, padding=0, bias=True)

    def forward(self, x):
        # 编码器路径
        conv_block_1 = self.conv_block_1(x) # 第一次卷积，保存用于跳跃连接
        max_pool_1 = self.max_pool(conv_block_1)

        conv_block_2 = self.conv_block_2(max_pool_1)  # 第二次卷积，保存用于跳跃连接
        max_pool_2 = self.max_pool(conv_block_2)

        conv_block_3 = self.conv_block_3(max_pool_2) # 第三次卷积，保存用于跳跃连接
        max_pool_3 = self.max_pool(conv_block_3)

        conv_block_4 = self.conv_block_4(max_pool_3) # 第四次卷积，保存用于跳跃连接
        max_pool_4 = self.max_pool(conv_block_4)
        # 瓶颈层
        conv_block_5 = self.conv_block_5(max_pool_4)
        # 解码器路径，包含跳跃连接
        up_block_1 = self.up_block_1(conv_block_5)
        conv_block_6 = self.conv_block_6(torch.cat([conv_block_4, up_block_1], dim=1)) # 跳跃连接

        up_block_2 = self.up_block_2(conv_block_6)
        conv_block_7 = self.conv_block_7(torch.cat([conv_block_3, up_block_2], dim=1)) # 跳跃连接

        up_block_3 = self.up_block_3(conv_block_7)
        conv_block_8 = self.conv_block_8(torch.cat([conv_block_2, up_block_3], dim=1)) # 跳跃连接

        up_block_4 = self.up_block_4(conv_block_8)
        conv_block_9 = self.conv_block_9(torch.cat([conv_block_1, up_block_4], dim=1)) # 跳跃连接
        # 最终输出层
        out = self.out(conv_block_9)

        return out

# 模型测试与参数统计
if __name__ == '__main__':
    device = torch.device("cpu")
    # 创建模型实例：输入通道1，输出通道14（类别数），初始特征数16
    model = UNet3D().to(device)
    # 打印模型摘要信息：批量大小8，输入尺寸为64×96×96的3D图像
    summary(model, (8, 1, 64, 96, 96))